package com.lktx.sso.config;

import cn.dev33.satoken.jwt.SaJwtTemplate;
import cn.dev33.satoken.jwt.SaJwtUtil;
import cn.dev33.satoken.jwt.StpLogicJwtForSimple;
import cn.dev33.satoken.stp.StpUtil;
import cn.hserver.core.interfaces.InitRunner;
import cn.hserver.core.ioc.annotation.Autowired;
import cn.hserver.core.ioc.annotation.Bean;
import cn.hserver.plugin.web.context.HServerContextHolder;
import cn.hutool.crypto.asymmetric.RSA;
import cn.hutool.jwt.JWT;
import cn.hutool.jwt.signers.JWTSigner;
import cn.hutool.jwt.signers.JWTSignerUtil;
import com.lktx.sso.admin.dto.KvDTO;
import com.lktx.sso.admin.entity.SsoAppJwtConfig;
import com.lktx.sso.admin.service.SsoAppJwtConfigService;
import lombok.extern.slf4j.Slf4j;

import java.util.HashMap;
import java.util.Map;

@Slf4j
@Bean
public class JwtConfig implements InitRunner {

    @Autowired
    private SsoAppJwtConfigService ssoAppJwtConfigService;

    @Override
    public void init(String[] args) {
        StpUtil.setStpLogic(new StpLogicJwtForSimple());
        SaJwtUtil.setSaJwtTemplate(new SaJwtTemplate() {
            @Override
            public String generateToken(JWT jwt, String keyt) {
                try {
                    String appId = HServerContextHolder.getWebKit().httpRequest.query("appId");
                    SsoAppJwtConfig byId = ssoAppJwtConfigService.getById(Integer.valueOf(appId));
                    log.info("------ 自定义了 token 生成算法: {}", byId.getSsoAppId());
                    if (byId.getJwtPayload() != null) {
                        for (KvDTO kvDTO : byId.getJwtPayload()) {
                            jwt.getPayload().setPayload(kvDTO.getKey(), kvDTO.getValue());
                        }
                    }
                    if (byId.getJwtHeader() != null) {
                        Map<String, Object> data = new HashMap<>();
                        for (KvDTO kvDTO : byId.getJwtHeader()) {
                            data.put(kvDTO.getKey(), kvDTO.getValue());
                        }
                        jwt.getHeader().addHeaders(data);
                    }
                    RSA rsa = new RSA(byId.getPrivateKey(), null);
                    JWTSigner jwtSigner;
                    if (byId.getIdTokenSignatureAlgorithm().equals("RS256")) {
                        jwtSigner = JWTSignerUtil.rs256(rsa.getPrivateKey());
                    } else if (byId.getIdTokenSignatureAlgorithm().equals("ES256")) {
                        jwtSigner = JWTSignerUtil.es256(rsa.getPrivateKey());
                    } else {
                        jwtSigner = JWTSignerUtil.hs256(byId.getPublicKey().getBytes());
                    }
                    return jwt.setSigner(jwtSigner).sign();
                } catch (Exception ignored) {
                }
                return super.generateToken(jwt, keyt);
            }
        });
    }
}
